#pragma once

#include "allocator.hpp"
#include "algo_meta_func.hpp"
#include "pair.hpp"
#include "rbtree.hpp"
#include "functor.hpp"

namespace zuck
{
	template<bool multi, typename Key, typename Val, typename Cmp, typename Equal, typename Alloc>
	class Map_Traits {
	private:
		using _Alty_tts = zuck::allocator_traits<Alloc>;
	public:
		using key_type		 = Key;
		using value_type	 = zuck::pair<Key, Val>;
		using equal			 = Equal;
		using key_compare	 = Cmp;
		using allocator_type = Alloc;

		inline static bool constexpr _Multi = multi;
	};

	template<
		typename K, typename V, 
		typename Cmp   = zuck::less<K>, 
		typename Equal = zuck::equal_to<K>, 
		typename Alloc = zuck::allocator<pair<K, V>>
	>
	class map : public _rbtree<Map_Traits<false, K, V, Cmp, Equal, Alloc>>
	{
		static_assert(
			!zuck::is_const_v<V> || !zuck::is_volatile_v<V>,
			"template parameter can not be const or volatile");
	private:
		using _Base     = _rbtree<Map_Traits<false, K, V, Cmp, Equal, Alloc>>;
		using _Self     = map<K, V, Cmp, Equal, Alloc>;
		using _Node_ptr = __tree_node<false, zuck::pair<K, V>>*;
	public:
		using key			   = K;
		using pointer		   = typename _Base::pointer;
		using value_type	   = typename _Base::value_type;
		using pointer		   = typename _Base::pointer;
		using size_type		   = typename _Base::size_type;
		using reference		   = typename _Base::reference;
		using const_reference  = typename _Base::const_reference;
		using rvalue_reference = typename _Base::rvalue_reference;
	public:
		using iterator				 = typename _Base::iterator;
		using const_iterator		 = typename _Base::const_iterator;
		using reverse_iterator		 = typename _Base::reverse_iterator;
		using cosnt_reverse_iterator = typename _Base::const_reverse_iterator;
	public:
		map() : _Base() {}

		map(std::initializer_list<value_type> data)
			: _Base(data) {}

        map(std::initializer_list<value_type> data, Alloc al)
			: _Base(data, al) {}

		template<
			typename InputIterator,
			zuck::enable_if_t<zuck::is_input_iterator_v<InputIterator>, int> = 0>
		map(InputIterator Begin, InputIterator End)
			: _Base(Begin, End) {}

		template<
			typename InputIterator,
			zuck::enable_if_t<zuck::is_input_iterator_v<InputIterator>, int> = 0>
		map(InputIterator Begin, InputIterator End, Alloc al)
			: _Base(Begin, End, al) {}

		map(const Cmp& _pr) : _Base(_pr) {}

        map(const Cmp& _Pr, Alloc al) : _Base(_Pr, al) {}

		map(const map& other, const Cmp& _pr) : _Base(other, _pr) {}

        map(const map& other, const Cmp& _Pr, Alloc al)
			: _Base(other, _Pr, al) {}

		map(const map& other) : _Base(other) {}

        map(const map& other, Alloc al) : _Base(other, al) {}

		map(map&& other)noexcept : _Base(zuck::move(other)) {}

		[[nodiscard]] V& at(const key& _key)
		{
			return __try_insert(_key);
		}

		[[nodiscard]] inline const V& at(key&& _key)const
		{
			return __try_insert(_key);
		}

		[[nodiscard]] size_type erase(const K& _k)
		{
			if (__find_node(_k) == iterator(nil<false, value_type>)) {
				return 0;
			}
			_Base::_m_root = __remove_node(_Base::_m_root, _k);
			_Base::_m_root->__col = _Color::black;
			return 1;
		}

		[[nodiscard]] size_type erase(K&& _k)
		{
			if (__find_node(_k) == iterator(nir<false, value_type>, _Base::_Max())) {
				return 0;
			}
			_Base::_m_root = __remove_node(_Base::_m_root, _k);
			_Base::_m_root->__col = _Color::black;
			return 1;
		}

		pair<iterator, bool> insert_or_assign(key const& _k, V const& _v)
		{
			return try_emplace(_k, _v);
		}

		pair<iterator, bool> insert_or_assign(key const& _k, V&& _v)
		{
			return try_emplace(_k, zuck::forward<V>(_v));
		}

		template<typename...Args>
		inline pair<iterator, bool> emplace(Args&&... args)
		{
			return __insert_aux(value_type(zuck::forward<Args>(args))...);
		}

		template<typename...Args>
		inline pair<iterator, bool> try_emplace(const key& _k, Args&&...args)
		{
			auto _Iter = find(_k);
			if (_Iter) {
				*(_Iter->__d) = V(zuck::forward(args)...);
				return { _Iter, true };
			}
			return __insert_aux(value_type(_k, V(zuck::forward(args)...)));
		}

		template<typename...Args>
		inline pair<iterator, bool> try_emplace(key&& _k, Args&&...args)
		{
			auto _Iter = find(_k);
			if (_Iter) {
				*(_Iter->__d) = V(zuck::forward(args)...);
				return { _Iter, true };
			}
			return __insert_aux(value_type(zuck::move(_k), 
				                V(zuck::forward(args)...)));
		}

		[[nodiscard]] inline V& operator[](const key& _key)
		{
			return __try_insert(_key);
		}

		[[nodiscard]] inline V& operator[](key&& _key)
		{
			return __try_insert(_key);
		}

		inline map& operator=(const map& other)
		{
			_Base::operator=(other);
			return *this;
		}

		inline map& operator=(map&& other)noexcept
		{
			_Base::operator=(zuck::move(other));
			return *this;
		}

	private:
		__tree_node<false, value_type>* __remove_node(_Node_ptr root, const K& val)
		{
			if (root == nil<false, value_type>) {
				return nil<false, value_type>; 
			}
			if (_Base::__cmp(val, root->__d->first)) {
				root->__l = __remove_node(root->__l, val);
			}
			else if (_Base::__cmp(root->__d.first, val)) {
				root->__r = __remove_node(root->__r, val);
			}
			else {
				if (root->__l == nil<false, value_type> || 
					root->__r == nil<false, value_type>) {
					_Node_ptr tmp = ((root->__l == nil<false, value_type>) ? 
										root->__r : root->__l);
					_Base::alloc.deallocate(root);
					--_Base::_m_cnt;
					return tmp;
				}
				else {
					_Node_ptr _pre = _Base::__predecessor(root);
					root->__d = _pre->__d;
					root->__l = __remove_node(root->__l, _pre->__d->first);
				}
			}
			return _Base::__remove_balance(root);
		}

		[[nodiscard]] inline V& __try_insert(const K& _key)
		{
			iterator iter = _Base::__find_node(_key);
			if (iter == iterator(nil<false, value_type>)) {
				return __insert_aux(value_type(_key, V{})).first->second;
			}
			return iter->second;
		}

		pair<iterator, bool> __insert_aux(const value_type& val)
		{
			_Base::_m_root = _Base::__insert(nir<false, value_type>, 
				                             _Base::_m_root, val);
			_Base::_m_root->__col = _Color::black;
			_Base::_m_root->_m_p  = nir<false, value_type>;
			nir<false, value_type>->__l = _Base::_m_root;

			return _Base::__exist ? 
				pair<iterator, bool>(iterator(nir<false, value_type>, 
					                 _Base::_Max()), !_Base::__exist) : 
				pair<iterator, bool>(_Base::__find_node(val.first), !_Base::__exist);
		}

	};

	template<typename K, typename V, 
		typename _Cmp   = zuck::less<K>, 
		typename _Equal = zuck::equal_to<K>, 
		typename Alloc = zuck::allocator<zuck::pair<K, V>>>
	map(std::initializer_list<zuck::pair<K, V>>)->map<K, V>;

	template<typename K, typename V, 
		typename _Cmp   = zuck::less<K>, 
		typename _Equal = zuck::equal_to<K>, 
		typename Alloc = zuck::allocator<zuck::pair<K, V>>>
	map(map<K, V> const&)->map<K, V, _Cmp, Alloc>;

	template<typename K, typename V,
		typename _Cmp   = zuck::less<K>,
		typename _Equal = zuck::equal_to<K>,
		typename Alloc = zuck::allocator<zuck::pair<K, V>>>
	map(map<K, V>&&)->map<K, V, _Cmp, Alloc>;

	template<typename K, typename V, 
		typename _Cmp = zuck::less<K>, 
		typename _Equal = zuck::equal_to<K>, 
		typename Alloc = zuck::allocator<zuck::pair<K, V>>>
	using ZMap = map<K, V, _Cmp, _Equal, Alloc>;
}

